# viz_demo.py
from pathlib import Path
import argparse
import webbrowser

from core.extractor import LineageExtractor
from viz.tree import render_table_tree
from viz.sankey import render_column_sankey, render_column_to_column_sankey


def main():
    ap = argparse.ArgumentParser(description="SQL 血缘可视化 Demo")
    ap.add_argument("--sql", default="examples/demo.sql", help="SQL 文件路径")
    ap.add_argument("--out", default="out", help="输出目录")
    ap.add_argument("--topn", type=int, default=30, help="表→列图每表最多展示的字段数")
    ap.add_argument("--width", default="1200px", help="画布宽度（如 1200px 或 100%）")
    ap.add_argument("--height", default="700px", help="画布高度（如 700px）")
    ap.add_argument("--open", action="store_true", help="渲染后自动在浏览器打开")
    args = ap.parse_args()

    sql_path = Path(args.sql)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 1) 解析 SQL
    sql = sql_path.read_text(encoding="utf-8")
    res = LineageExtractor().extract(sql)
    if not res.statements:
        print("No statements parsed from SQL.")
        return
    st = res.statements[0]  # 仅演示第一条

    base = sql_path.stem

    # 2) 表血缘（树）
    tree_chart = render_table_tree(st, return_html=False, title="表血缘")
    out_tree = out_dir / f"{base}_table_tree.html"
    tree_chart.render(str(out_tree))

    # 3) 字段血缘（表→列）Sankey
    sankey_chart = render_column_sankey(
        st,
        return_html=False,
        title="字段血缘（表 → 列）",
        top_n_columns=args.topn,
        width=args.width,
        height=args.height,
    )
    out_col = out_dir / f"{base}_column_sankey.html"
    sankey_chart.render(str(out_col))

    # 4) 列到列（源列 → 目标列/别名）Sankey
    col2col_chart = render_column_to_column_sankey(
        st,
        return_html=False,
        title="字段血缘（源列 → 目标列）",
        width=args.width,
        height=args.height,
    )
    out_c2c = out_dir / f"{base}_col2col_sankey.html"
    col2col_chart.render(str(out_c2c))

    print("Done:")
    print(f"  {out_tree}")
    print(f"  {out_col}")
    print(f"  {out_c2c}")

    if args.open:
        try:
            # 先转绝对路径，再转 file:// URI
            webbrowser.open(out_tree.resolve().as_uri())
        except Exception:
            # 极端环境兜底：直接打开绝对路径
            import sys, os
            p = str(out_tree.resolve())
            if sys.platform.startswith("win") and hasattr(os, "startfile"):
                os.startfile(p)  # Windows 原生打开
            else:
                webbrowser.open_new_tab(p)


if __name__ == "__main__":
    main()
